import json
import os

from OCC.Core.BRepBuilderAPI import BRepBuilderAPI_Transform
from OCC.Core.IFSelect import IFSelect_RetDone, IFSelect_ItemsByEntity
from OCC.Core.Quantity import Quantity_Color, Quantity_TOC_RGB
from OCC.Core.STEPCAFControl import STEPCAFControl_Reader
from OCC.Core.STEPControl import STEPControl_Reader
from OCC.Core.TCollection import TCollection_ExtendedString
from OCC.Core.TDF import TDF_LabelSequence, TDF_Label
from OCC.Core.TDocStd import TDocStd_Document
from OCC.Core.Tesselator import ShapeTesselator
from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.XCAFDoc import XCAFDoc_DocumentTool_ShapeTool, XCAFDoc_DocumentTool_ColorTool


def read_step_file_with_names_colors_(filename):
    """ Returns list of tuples (topods_shape, label, color)
    Use OCAF.
    """
    if not os.path.isfile(filename):
        raise FileNotFoundError("%s not found." % filename)
    # the list:
    output_shapes = {}
    color_list = []
    # create an handle to a document
    doc = TDocStd_Document(TCollection_ExtendedString("pythonocc-doc"))

    # Get root assembly
    shape_tool = XCAFDoc_DocumentTool_ShapeTool(doc.Main())
    color_tool = XCAFDoc_DocumentTool_ColorTool(doc.Main())
    # layer_tool = XCAFDoc_DocumentTool_LayerTool(doc.Main())
    # mat_tool = XCAFDoc_DocumentTool_MaterialTool(doc.Main())

    step_reader = STEPCAFControl_Reader()
    step_reader.SetColorMode(True)
    step_reader.SetLayerMode(True)
    step_reader.SetNameMode(True)
    step_reader.SetMatMode(True)
    step_reader.SetGDTMode(True)

    status = step_reader.ReadFile(filename)
    if status == IFSelect_RetDone:
        step_reader.Transfer(doc)

    locs = []

    def _get_sub_shapes(lab, loc):
        # global cnt, lvl
        # cnt += 1
        # users = TDF_LabelSequence()
        # users_cnt = shape_tool.GetUsers(lab, users)

        l_subss = TDF_LabelSequence()
        shape_tool.GetSubShapes(lab, l_subss)

        l_comps = TDF_LabelSequence()
        shape_tool.GetComponents(lab, l_comps)

        name = lab.GetLabelName()
        print("Name :", name)

        if shape_tool.IsAssembly(lab):
            l_c = TDF_LabelSequence()
            shape_tool.GetComponents(lab, l_c)
            for i in range(l_c.Length()):
                label = l_c.Value(i + 1)
                if shape_tool.IsReference(label):
                    label_reference = TDF_Label()
                    shape_tool.GetReferredShape(label, label_reference)
                    loc = shape_tool.GetLocation(label)
                    # trans = loc.Transformation()
                    # rot = trans.GetRotation()
                    # tran = trans.TranslationPart()

                    locs.append(loc)
                    # lvl += 1
                    _get_sub_shapes(label_reference, loc)
                    # lvl -= 1
                    locs.pop()

        elif shape_tool.IsSimpleShape(lab):
            shape = shape_tool.GetShape(lab)
            print(1111111111111111111111111111111111111111111111, shape)
            loc = TopLoc_Location()
            for l in locs:
                loc = loc.Multiplied(l)

            # trans = loc.Transformation()
            # rot = trans.GetRotation()
            # tran = trans.TranslationPart()
            c = Quantity_Color(0.5, 0.5, 0.5, Quantity_TOC_RGB)  # default color
            colorSet = False
            if (color_tool.GetInstanceColor(shape, 0, c) or
                    color_tool.GetInstanceColor(shape, 1, c) or
                    color_tool.GetInstanceColor(shape, 2, c)):
                color_tool.SetInstanceColor(shape, 0, c)
                color_tool.SetInstanceColor(shape, 1, c)
                color_tool.SetInstanceColor(shape, 2, c)
                colorSet = True
                n = c.Name(c.Red(), c.Green(), c.Blue())
                color_list.extend([c.Red(), c.Green(), c.Blue()])
            if not colorSet:
                if (color_tool.GetColor(lab, 0, c) or
                        color_tool.GetColor(lab, 1, c) or
                        color_tool.GetColor(lab, 2, c)):
                    color_tool.SetInstanceColor(shape, 0, c)
                    color_tool.SetInstanceColor(shape, 1, c)
                    color_tool.SetInstanceColor(shape, 2, c)

                    n = c.Name(c.Red(), c.Green(), c.Blue())
                    color_list.extend([c.Red(), c.Green(), c.Blue()])
            shape_disp = BRepBuilderAPI_Transform(shape, loc.Transformation()).Shape()
            if not shape_disp in output_shapes:
                output_shapes[shape_disp] = [lab.GetLabelName(), c]
            for i in range(l_subss.Length()):
                lab_subs = l_subss.Value(i + 1)
                shape_sub = shape_tool.GetShape(lab_subs)

                c = Quantity_Color(0.5, 0.5, 0.5, Quantity_TOC_RGB)  # default color
                colorSet = False
                if (color_tool.GetInstanceColor(shape_sub, 0, c) or
                        color_tool.GetInstanceColor(shape_sub, 1, c) or
                        color_tool.GetInstanceColor(shape_sub, 2, c)):
                    color_tool.SetInstanceColor(shape_sub, 0, c)
                    color_tool.SetInstanceColor(shape_sub, 1, c)
                    color_tool.SetInstanceColor(shape_sub, 2, c)
                    colorSet = True
                    n = c.Name(c.Red(), c.Green(), c.Blue())
                    color_list.extend([c.Red(), c.Green(), c.Blue()])
                if not colorSet:
                    if (color_tool.GetColor(lab_subs, 0, c) or
                            color_tool.GetColor(lab_subs, 1, c) or
                            color_tool.GetColor(lab_subs, 2, c)):
                        color_tool.SetInstanceColor(shape, 0, c)
                        color_tool.SetInstanceColor(shape, 1, c)
                        color_tool.SetInstanceColor(shape, 2, c)

                        n = c.Name(c.Red(), c.Green(), c.Blue())
                        color_list.extend([c.Red(), c.Green(), c.Blue()])
                shape_to_disp = BRepBuilderAPI_Transform(shape_sub, loc.Transformation()).Shape()
                # position the subshape to display
                if not shape_to_disp in output_shapes:
                    output_shapes[shape_to_disp] = [lab_subs.GetLabelName(), c]


    def _get_shapes():
        labels = TDF_LabelSequence()
        shape_tool.GetFreeShapes(labels)
        # global cnt
        # cnt += 1

        print()
        print("Number of shapes at root:", labels.Length())
        print()
        for i in range(labels.Length()):
            root_item = labels.Value(i + 1)
            print(root_item)
            _get_sub_shapes(root_item, None)

    _get_shapes()

    step_reader_ = STEPControl_Reader()
    status = step_reader_.ReadFile(filename)
    if status == IFSelect_RetDone:  # check status
        failsonly = False
        step_reader_.PrintCheckLoad(failsonly, IFSelect_ItemsByEntity)
        step_reader_.PrintCheckTransfer(failsonly, IFSelect_ItemsByEntity)

        step_reader_.TransferRoot(1)
        _nbs = step_reader_.NbShapes()
        _shape = step_reader_.Shape(1)
        _tess = ShapeTesselator(_shape)
        _tess.Compute(compute_edges=False, mesh_quality=50)

        with open("bbb" + ".json", "w") as text_file:
            json_ = _tess.ExportShapeToThreejsJSONString(filename)
            print("33333333333333", json)
            json_ = json_.replace("\\", "/")
            json_d = json.loads(json_)
            json_d["data"]["attributes"]["color"] = {}
            json_d["data"]["attributes"]["color"]["itemSize"] = 3
            json_d["data"]["attributes"]["color"]["type"] = "Float32Array"
            json_d["data"]["attributes"]["color"]["array"] = color_list
            text_file.write(json.dumps(json_d))
    else:
        raise Exception("Error: can't read file - Method: _load_STEP_file")
    return output_shapes
